# -----------------------------------------------------------------------------
# Copyright (c) 2025, Southeast University (China)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name of the copyright holders nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Jiajie Xu, Yanfang Zhang, Jiaqi Gao, Leyun Tian
# -----------------------------------------------------------------------------


from Data_var import *
import work_var
import networkx as nx

class TimingGraph:
    """Timing Graph for test design"""
    """
    self.design -> current design name
    self.cells -> a name dict for cells: cell refname (U222) -> cell name (ND2D0BWP16P90)
    self.design_max_trans -> the max transition time of design
    self.design_max_caps -> the max capacitance of design
    self.G -> DAG timing graphs
            G -> nodes: named with pin/port names
            features:
            'arrival': arrival time for pin and ports: [Max_Rise, Max_Fall, Min_Rise, Min_Fall]
            'ceff': effective capacitance calculated by PrimeTime for each driver pin: [Max_Rise_Ceff, Max_Fall_Ceff, Min_Rise_Ceff, Min_Fall_Ceff]
            'slack': slack for pin and ports: [Max_Rise, Max_Fall, Min_Rise, Min_Fall]
            'trans': transition time for pin and ports: [Max_Rise, Max_Fall, Min_Rise, Min_Fall]
            'is_port': if the node is a PI/PO: True/False
            'max_trans': maxium transition time by DRV
            'max_caps': maxium capacitance set by DRV
            'is_async_pin': if the pin is an asynchronous preset/clear pin
            'is_clock_network': if the pin is a combinational fanout of a clock
            
            G -> edges features:
            'name': cell/net name for different arcs
            'leakage': PT total leakage for cell arcs, None for net arcs
            'is_cell': if the edge is a cell arc
            'delay': the delay for the timing arc: [Max_Rise, Max_Fall, Min_Rise, Min_Fall]
            'sense_unate': the unate of the timing sense
            'when': when condition of the timing arc
            'sdf_cond': sdf condition of the timing arc
    
    function:
    print_nodes(self): output pin/port features for current design
    print_edges(self): output cell/net arc features for current design
    """
    
    def __init__(self, design):
        print(f"Initializing DAG timing graph of design {design} from PrimeTime Rpt")
        """Initialize DAG timing graph"""
        self.design = design
        
        # define cell and net Rpt Path
        CellRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_cell.rpt")
        NetRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_net.rpt")
        # cells with name, type, inpins and outpins
        cells = PtCellRpt_Parser.Read_PtCellRpt(CellRpt_Path)
        # nets with name inpins outpins resistance and capacitance
        nets = PtNetRpt_Parser.Read_PtNetRpt(NetRpt_Path)
        
        # define the PT timing arc Rpt Path
        TimingArcRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_timing_arcs.rpt")
        timing_arcs = PtTimingArc_Parser.Read_PtTimingArc(TimingArcRpt_Path)
        
        # refname dict
        self.cells = {}
        for cell in cells.values():
            self.cells[cell.name] = cell.type # U222 -> ND2D0BWP16P90 refname -> libname
        
        # Build Graph from Rpt items
        """ 
        Timing Graph G: multi edge DAG
        """
        self.G = nx.MultiDiGraph()
        """Noted, the delay is annotated in Max_Rise, Max_Fall, Min_Rise, Min_Fall order"""
        for timing_arc in timing_arcs: 
            edge_data = {
                "is_cell": timing_arc.is_cell,
                "delay": timing_arc.delay,
                "sense_unate": timing_arc.sense_unate,
                "sdf_cond": timing_arc.sdf_cond,
                "when": timing_arc.when,
            }

            if timing_arc.is_cell:
                edge_data["name"] = timing_arc.from_pin.split("/")[0]

            self.G.add_edge(timing_arc.from_pin, timing_arc.to_pin, **edge_data)
        
        for net_name, net in nets.items():
            for inp in net.inpins:
                for outp in net.outpins:  
                    for key in self.G[inp][outp]:
                        self.G[inp][outp][key]["name"] = net_name
        
        """Reading data from PrimeTime"""
        # define Pt Data Output Rpt Path
        ClockPinRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/func_slow_data_used_as_clocks.txt.gz")
        # pin arrival time
        PinArrivalRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_pin_arrival.rpt")
        # port arrival time
        PortArrivalRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_port_arrival.rpt")
        # pin effective capacitance
        PinCapsRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_pin_caps.rpt")
        # slack
        PinSlackRpt_Path = os.path.join(work_var.pt_data_dir, f"{design}/func_slow_data_pin_slack.txt.gz")
        # pin transition time
        PinTrans_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_pin_transition.rpt")
        # max transition time and max pin transition time and port transition time
        PinTiming_Path = os.path.join(work_var.pt_data_dir, f"{design}/func_slow_data_pin_timing.txt.gz")
        # cell leakage power
        CellLeakage_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_leakage.rpt")
        # pin type
        PinType_Path = os.path.join(work_var.pt_data_dir, f"{design}/{design}_pin_type.rpt")
        
        clock_pins, _ = PtClockPin_Parser.Read_PtClockPin(ClockPinRpt_Path)
        # pins -> Max_Rise, Max_Fall, Min_Rise, Min_Fall arrival
        pins_arrival = PtPinArrival_Parser.Read_PtPinArrival(PinArrivalRpt_Path)
        # ports and Max_Rise, Max_Fall, Min_Rise, Min_Fall arrival
        ports_arrival = PtPortArrival_Parser.Read_PtPortArrival(PortArrivalRpt_Path)
        # pin and Max_Rise, Max_Fall, Min_Rise, Min_Fall effective capacitance
        pins_caps = PtPinCaps_Parser.Read_PtPinCaps(PinCapsRpt_Path)
        # pins -> Max_Rise, Max_Fall, Min_Rise, Min_Fall slack
        pins_ports_slack, self.timing_unit = PtPinSlack_Parser.Read_PtPinSlack(PinSlackRpt_Path)
        # pins -> Max_Rise, Max_Fall, Min_Rise, Min_Fall transition
        pins_trans = PtPinActualTrans_Parser.Read_PtPinActualTrans(PinTrans_Path)
        # pins -> Max_Rise Transition, ports and SetupRise, SetupFall, HoldRise, HoldFall, MaxTrans, MaxCap data
        self.design_max_trans, self.design_max_caps, pins_maxtrans, ports_data = PtPinTiming_Parser.Read_PtTrans(PinTiming_Path)
        # cells -> Gate_Leakage, Intrinsic_Leakage,Totgal_Leakage
        cells_leakage = PtLeakage_Parser.Read_PtLeakage(CellLeakage_Path)
        # pins -> is_async_pin
        pin_types = PtPinType_Parser.Read_PtPinType(PinType_Path)
        
        """ annotate timing features on pins"""
        """ noted the timing parameters are annotated in Max_Rise, Max_Fall, Min_Rise, Min_Fall order"""
        for node in self.G.nodes:
            # if node is used as a clock pin
            if node in clock_pins:
                self.G.nodes[node].update({'is_clock_pin': True})
            else:
                self.G.nodes[node].update({'is_clock_pin': False})
            # if node is a port
            if node in ports_data:
                port_data = ports_data[node]
                self.G.nodes[node].update({
                    'is_port': True,
                    # SetupRise -> Max_Rise, SetupFall -> Max_Fall, HoldRise -> Min_Rise, HoldFall -> Min_Fall
                    'trans': [port_data.data[0], port_data.data[1], port_data.data[2], port_data.data[3]],
                    'max_trans': port_data.data[4],
                    'max_caps': port_data.data[5],
                    'ceff': [None, None, None, None]
                })
                # arrival
                port_arrival = ports_arrival.get(node)
                self.G.nodes[node]['arrival'] = port_arrival.arrival if port_arrival else [None, None, None, None]
            
            # if the node is a pin
            else:
                self.G.nodes[node]['is_port'] = False
                # arrival time
                pin_arrival = pins_arrival.get(node)
                self.G.nodes[node]['arrival'] = pin_arrival.arrival if pin_arrival else [None, None, None, None]
                # effective capacitance 
                pin_caps = pins_caps.get(node)
                self.G.nodes[node]['ceff'] = pin_caps.ceff if pin_caps else [None, None, None, None]
                # transition time
                pin_trans = pins_trans.get(node)
                self.G.nodes[node]['trans'] = pin_trans.trans if pin_trans else [None, None, None, None]
                # max_transition time
                pin_maxtrans = pins_maxtrans.get(node)
                self.G.nodes[node]['max_trans'] = pin_maxtrans.trans[0] if pin_maxtrans else None
                # max_capacitance
                self.G.nodes[node]['max_caps'] = None

            # slack
            pin_port_slack = pins_ports_slack.get(node)
            self.G.nodes[node]['slack'] = pin_port_slack.slack if pin_port_slack else [None, None, None, None]
            
            # type
            type = pin_types[node]
            self.G.nodes[node]['is_async_pin'] = type[0]
            self.G.nodes[node]['is_clock_network'] = type[1]
        
        # annotated leakage power
        for u, v, key, data in self.G.edges(keys=True, data=True):
            if data.get('is_cell'):
                self.G[u][v][key]["leakage"] = cells_leakage[u.split('/')[0]].leakage[2]
            else:
                self.G[u][v][key]["leakage"] = None
        
        print(f"Finished Initializing Timing Graph")
    
    def print_nodes(self, graph = None, width=30):
        """
        print node name and node features in self.G
        """
        if not graph:
            Graph = self.G
        else:
            Graph = graph
        print("------------------------------------------ Node Information ------------------------------------------")
        for node, data in Graph.nodes(data=True):
            node_str = str(node).ljust(width)
            print(f"Node: {node_str} | Data: {data}")
        print("------------------------------------------------------------------------------------------------------\n")
        
    def print_edges(self, graph = None, width=30):
        """
        print edge name and edge features in self.G
        """
        if not graph:
            Graph = self.G
        else:
            Graph = graph
        print("------------------------------------------------------------------------------------ Edge Information ------------------------------------------------------------------------------------")
        for u, v, key, data in Graph.edges(keys=True, data=True):
            u_str = str(u).ljust(width)
            v_str = str(v).ljust(width)
            print(f"Edge: {u_str} -> {v_str} :edge{key} | Data: {data}")
        print("------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n")
    
if __name__ == "__main__": 
    TG = TimingGraph("mc_top")
    TG.print_nodes()
    pass